

#include <paging.h>
#include <pmm.h>
#include <string.h> 
#include <msr.h>
#include <debug.h>
#include <process.h>

static dword *pd;

int pg_map_frames(dword *pd, void *start, void *end);

/**
 * Initialize paging mechanism
 * Identity maps the kernel space.
 * 
 * returns - error
 */
int init_paging(void)
{
	int ret = 0;
	
	pd = (dword*)pmm_frame_alloc(1);
	bzero((void*)pd, 4096);
	
	ret = pg_map_frames(pd, 0, (void*)g_highmem);
	
	//TODO: map kernel pages with GLOBAL flag set
	/*
	while (ptr < g_highmem)
	{
		//pg_map_frame((void*)ptr, pd);
		
		if ((ptr & 0x3FFFFF) == 0)
		{
			// begins a 4MB region, alloc new PDE
			pde = pg_pagetable_alloc(1, 0, 0, 0, 0);  // allocate pagetable
			pd[next_pde] = pde;  // set entry in pd
			ppde = (dword*)(pde & ~0xFFF);  // get pde address
			next_pte = 0;
			
			// previous PDE must be full
			if (next_pde > 0)  // not the first pde
				pd[next_pde-1] = pd[next_pde-1] | (1 << PG_PDE_FULL);
				
			next_pde++;
		}
		
		if ((ptr & 0xFFF) == 0)
		{
			// begins a 4K region, alloc new PTE
			pte = pg_pte_alloc(ptr, 1, 0, 0, 0, 0);  // create PTE
			ppde[next_pte] = pte;  // set PTE in PD
			next_pte++;
		}
		
		ptr += 4096;  // add 4K
	}
	*/
	// set the system process's pd
	proc_get_sysprocess()->pd = pd;
	
	msr_write_cr3((dword)pd);
	msr_write_cr0(msr_read_cr0() | 0x80000000);
	
	return ret;
}

/**
 * Gets the PDE and PTE for a given address
 * 
 * void *addr - virtual address
 * 
 * returns -
 *   dword *pde - ptr to PDE
 *   dword *pte - ptr to PTE
 */
void pg_get_structs_from_addr(void *addr, dword *pde, dword *pte)
{
	*pde = pd[((dword)addr & 0xFFC00000) >> 22];
	dword *ppde = (dword*)(*pde & ~0xFFF);
	*pte = ppde[((dword)addr & 0x003FF000) >> 12];
}

/**
 * Allocates a pageteble entry with specified attributes
 *
 * dword addr - address to map to 
 * byte rw - read/write
 * byte us - user/supervisor
 * byte wtr - write-through
 * byte pcd - cache disabled
 * byte g - global
 * 
 * returns - pagetable entry, NULL on error
 */
dword pg_pte_alloc(dword addr, byte rw, byte us, byte wtr, byte pcd, byte g)
{
	dword pte = 0;
	
	pte = (addr | (rw << PG_PDE_RW) | (us << PG_PDE_US) | (pcd << PG_PDE_PCD) |
			(1 << PG_PDE_P) | (wtr << PG_PDE_PWT) | (g << PG_PDE_GLOBAL));
	
	return pte;
}

/**
 * Allocates a pagetable with specified attributes
 * 
 * byte rw - read/write
 * byte us - user/supervisor
 * byte wtr - write-through
 * byte pcd - cache disabled
 * byte g - global
 * 
 * returns - pagedirectory entry, NULL on error
 */
dword pg_pagetable_alloc(byte rw, byte us, byte wtr, byte pcd, byte g)
{
	dword pde = 0;

	dword *addr;
	addr = pmm_frame_alloc(1);
	if (addr != NULL)
	{
		bzero((void*)addr, 4096);
		pde = ((dword)addr | (rw << PG_PDE_RW) | (us << PG_PDE_US) | (pcd << PG_PDE_PCD) |
				(1 << PG_PDE_P) | (wtr << PG_PDE_PWT) | (g << PG_PDE_GLOBAL));
	}
	else
	{
		debugf("[PG] Could not allocate page table\n");
		pause();
		return NULL;
	}
	
	//debugf("[PG] Allocated 0x%08x\n", addr);
	
	return pde;
}

/**
 * Map a frame into a process's memory space
 * 
 * void *frame - frame to map
 * dword *pd - pagedirectory to map frame into
 * 
 * returns - virtual address of the new page relative to the given pd
 */
void * pg_map_frame(void *frame, dword *pd)
{
	void *virt = NULL;
	bool done = FALSE;
	dword pde;
	dword *ppde;
	dword pte;
	
	//debugf("[PG] Looking for spot for 0x%08x...", frame);
	
	// find a free page
	int i = 0, j;
	while ((!done) && (i < 1024))
	{
		pde = pd[i];
		if (((pde & (1 << PG_PDE_FULL)) == 0) && ((pde & (1 << PG_PDE_P)) != 0))  // not full, but present
		{
			// check each PTE in this PDE
			ppde = (dword*)(pde & ~0xFFF);
			j = 0;
			while ((!done) && (j < 1024))
			{
				pte = ppde[j];
				if ((pte & (1 << PG_PDE_P)) == 0)  // not present
				{
					ppde[j] = pg_pte_alloc((dword)frame, 1, 0, 0, 0, 0);
					virt = (void*)(((i & 0x3FF) << 22) | ((j & 0x3FF) << 12));
					
					// check if pde is full?
					bool full = TRUE;
					int k;
					for (k = 0; k < 1024; k++)
					{
						if ((ppde[k] & (1 << PG_PDE_P)) == 0)
						{
							full = FALSE;
							break;
						}
					}
					
					if (full == TRUE)
					{
						pde = pde | (1 << PG_PDE_FULL);
						pd[i] = pde;
					}
					
					done = TRUE;
				}
				j++;
			}
		}
		i++;
	}
	
	if ((done == FALSE) && (virt == NULL))
	{
		// allocate new pde
		dword newpde = pg_pagetable_alloc(1, 0, 0, 0, 0);
		
		// find free pde
		i = 0;
		while ((!done) && (i < 1024))
		{
			pde = pd[i];
			if ((pde & (1 << PG_PDE_P)) == 0)  // not present
			{
				pd[i] = newpde;
				
				// use first entry
				ppde = (dword*)(newpde & ~0xFFF);
				ppde[0] = pg_pte_alloc((dword)frame, 1, 0, 0, 0, 0);
				
				virt = (void*)((i & 0x3FF) << 22);  // pte = 0
				
				done = TRUE;
			}
			i++;
		}
		
		if ((done == FALSE) && (virt == NULL))
		{
			debugf("[PG] Could not find free PTE to map frame\n");
			pause();
		}
	}
	
	//debugf("[PG] Allocated 0x%08x\n", virt);
	
	return virt;
}

/**
 * Maps a block of frames into the begining of a pagedirectory
 * 
 * dword *pd - pagedirectory (empty)
 * void *start - starting address
 * void *end - ending address
 * 
 * returns - error
 */
int pg_map_frames(dword *pd, void *start, void *end)
{
	dword ptr = (dword)start;
	dword next_pde = 0;
	dword next_pte = 0;
	dword pde = 0;
	dword pte = 0;
	dword *ppde = 0;
	
	while (ptr < (dword)end)
	{
		if ((ptr & 0x3FFFFF) == 0)
		{
			// begins a 4MB region, alloc new PDE
			pde = pg_pagetable_alloc(1, 0, 0, 0, 0);  // allocate pagetable
			pd[next_pde] = pde;  // set entry in pd
			ppde = (dword*)(pde & ~0xFFF);  // get pde address
			next_pte = 0;
			
			// previous PDE must be full
			if (next_pde > 0)  // not the first pde
				pd[next_pde-1] = pd[next_pde-1] | (1 << PG_PDE_FULL);
				
			next_pde++;
		}
		
		if ((ptr & 0xFFF) == 0)
		{
			// begins a 4K region, alloc new PTE
			pte = pg_pte_alloc(ptr, 1, 0, 0, 0, 0);  // create PTE
			ppde[next_pte] = pte;  // set PTE in PD
			next_pte++;
		}
		
		ptr += 4096;
	}
	
	return 0;
}

/**
 * Creates an address space usable for a user process
 * Maps the kernel into a new address space
 * 
 * returns - ptr to pagedirectory, or NULL on error
 */
dword * pg_create_address_space(void)
{
	dword *pd;
	int ret;
	
	pd = pmm_frame_alloc(1);
	if (pd == NULL)
	{
		debugf("[PG] Could not allocate pagedirectory structure\n");
		return NULL;
	}
	
	bzero((void*)pd, 4096);
	
	// fill in entries for kernel
	ret = pg_map_frames(pd, 0, aligned_kend);
	
	return pd;
}

/**
 * Converts a physical address to a virtual address according to the given pagedirectory
 * 
 * void *addr - physical address
 * dword *pd - pagedirectory
 * 
 * returns - virtual address, or NULL on error
 */
dword pg_pa2va(void *addr, dword *pd)
{
	//((((dword*)(pd[((addr) & 0xFFC00000) >> 22])[((addr) & 0x003FF000) >> 12]) & ~0xFFF) | ((addr) & 0xFFF))
	dword pde = pd[((dword)addr & 0xFFC00000) >> 22];
	dword *ppde = (dword*)(pde & ~0xFFF);
	dword pte = ppde[((dword)addr & 0x003FF000) >> 12];
	return ((pte & ~0xFFF) | ((dword)addr & 0xFFF));
}
